Greptile Summary

This PR adds unswizzle support for MXFP8 / NVFP4 scaling factors, implementing the inverse of the existing swizzle operation. It introduces new CUDA kernels (unswizzle_row_scaling_kernel_impl, unswizzle_col_scaling_kernel_impl), single-tensor and multi-tensor host-side dispatch functions, public C API declarations, and a comprehensive test suite including padded-shape standalone unswizzle tests and swizzle→unswizzle round-trip tests.

• The kernel logic is a correct structural inverse of the swizzle kernels: shared-memory roles are swapped (linear load from swizzled global → SLM, then read via swizzle index vs. write via swizzle index → linear store to swizzled global), and regs_unshuffle / regs_unshuffle_with_bit_shifts are the correct inverses of their regs_shuffle* counterparts.
• Output-size validation in unswizzle_scaling_factors correctly uses padded m * k from the swizzled input's shape, aligning with the multi-tensor path and the swizzle counterpart.
• The test suite uses padded_dim0 * padded_dim1 for the unswizzle comparison (covering padding bytes), and the unswizzle data shapes include several padded-boundary cases (M-only, K-only, both).
• One logic issue: unswizzle_scaling_factors (and multi_tensor_unswizzle_scaling_factors) rejects input tensors that carry both rowwise and columnwise scaling factors, but swizzle_scaling_factors processes both scale types in a single call. This asymmetry breaks the round-trip invariant for dual-scale MXFP8 tensors and is undocumented in the public header.

Confidence Score: 3/5

• Mostly safe to merge, but the asymmetric dual-scale rejection should be resolved before this lands in a release used for round-trip workflows.
• The kernel math and memory-layout inversions are correct and well-tested for single-scale tensors. The padding-aware test shapes cover the previously-reported validation bug. The remaining concern is an API contract issue: nvte_unswizzle_scaling_factors cannot round-trip a tensor that was produced by nvte_swizzle_scaling_factors if that tensor had both rowwise and columnwise scales — a configuration that is valid for MXFP8 dual-path training. This does not affect correctness for single-scale use cases, but it is a latent bug for callers who use the full MXFP8 dual-scale path.
• transformer_engine/common/swizzle/swizzle.cu — the dual-scale rejection check in unswizzle_scaling_factors (lines 1163–1166) and the corresponding check in multi_tensor_unswizzle_scaling_factors (lines 1391–1392).

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant Caller
    participant nvte_unswizzle_scaling_factors
    participant unswizzle_scaling_factors
    participant unswizzle_row_scaling_kernel
    participant unswizzle_col_scaling_kernel

    Caller->>nvte_unswizzle_scaling_factors: (swizzled_tensor, output_tensor, stream)
    nvte_unswizzle_scaling_factors->>unswizzle_scaling_factors: dispatch
    unswizzle_scaling_factors->>unswizzle_scaling_factors: validate scaling_mode, swizzled flag,\nm%128==0, k%4==0, output size == m*k
    alt rowwise_unswizzle
        unswizzle_scaling_factors->>unswizzle_row_scaling_kernel: <<<(DIVUP(tiles_k,n_tb), tiles_m)>>>\n(swizzled_ptr, compact_ptr, m, k)
        note over unswizzle_row_scaling_kernel: 1) Linear load: swizzled global → SLM\n2) __syncthreads\n3) SLM tile → regs (swizzle index)\n4) regs_unshuffle<LType>\n5) Write regs → compact global
    else columnwise_unswizzle
        unswizzle_scaling_factors->>unswizzle_col_scaling_kernel: <<<(DIVUP(tiles_k,TB_DIM), DIVUP(tiles_m,vls))>>>\n(swizzled_ptr, compact_ptr, m, k)
        note over unswizzle_col_scaling_kernel: 1) Linear load: swizzled global → SLM\n2) __syncthreads\n3) SLM (swizzle index) → regs\n4) regs_unshuffle_with_bit_shifts\n5) Write regs → compact global (M-major)
    end
    unswizzle_scaling_factors-->>Caller: compact scale_inv written to output

_{Last reviewed commit: 4410e9d}

@int-smart Please address the comments from Greptile and ideally also add the test case with the input not already padded to 128,128.

@ptrendx Will look into these

@ptrendx From what I am understanding then, there is no relevance of padding to the unswizzle kernel. Since the padding is already done during the swizzling operation I can just mirror it back to compact layout with the zero pads correctly in the compact layout and that should do. Is that assumption correct. Initially I was thinking of removing the padding from the scale_inv itself since this would be used for checkpointing